#------------------------------------------------#
#       evaluate 工具函数：按类别求mIoU, mPA
#------------------------------------------------#
import argparse
import json
from os.path import join

import numpy as np
from PIL import Image

'''
    # 内容：
        def fast_hist()     计算一张图片的标签分割图与预测分割图的混淆矩阵
        def per_class_iu()  按类求IoU
        def per_class_PA()  按类求？？
      > def compute_mIoU()  计算测试集的mIoU
        过程代码

    # 备注：
        mIoU的计算过程：
            1.计算每张图的IoU，使用混淆矩阵统计每个像素点的分类结果（有了混淆矩阵，就能算出这张图片的IoU = TP/(FN+FP)
            2.计算测试集上所有图片的混淆矩阵的直接加和（有了测试集混淆矩阵，就能算出每个类的mIoU_class = TP_total / (FN_total + FP_total)）
            3.所有分类的mIoU_class，直接平均，就得到了在这个测试集上的mIoU结果（nIoU = Σ mIoU_class / num_classes）
'''

# 设标签宽W，长H
def fast_hist(a, b, n):
    '''
        inputs:
            a   np.array.flatten()  真实标签
            b   np.array.flatten()  预测结果
            n   int                 分类的类别数量
    '''
    #--------------------------------------------------------------------------------#
    #   a是转化成一维数组的标签，形状(H×W,)；b是转化成一维数组的预测结果，形状(H×W,)
    #--------------------------------------------------------------------------------#
    k = (a >= 0) & (a < n) # 有效分类classIndex的掩码，np.array 一维 //  0 ~ num_class-1 为VOC数据集+背景类的分类索引
    #--------------------------------------------------------------------------------#
    #   np.bincount计算了从0到n**2-1这n**2个数中每个数出现的次数，返回值形状(n, n)
    #   返回中，写对角线上的为分类正确的像素点
    #--------------------------------------------------------------------------------#
    return np.bincount(n * a[k].astype(int) + b[k], minlength=n ** 2).reshape(n, n)  # 数学技巧？！！！！！！！！！！！！

def per_class_iu(hist):
    return np.diag(hist) / np.maximum((hist.sum(1) + hist.sum(0) - np.diag(hist)), 1) 

def per_class_PA(hist):
    return np.diag(hist) / np.maximum(hist.sum(1), 1) 

def compute_mIoU(gt_dir, pred_dir, png_name_list, num_classes, name_classes):  
    '''
        inputs:
            gt_dir          str     文件路径：真实标签图
            pred_dir        str     文件路径：预测灰度图
            png_name_list   list    样本图片的文件名列表
            num_classes     int     分类类的数量 + 背景
            name_classes    list    分类的类别名称
        outputs:
            mIoUs
        关键的过程变量：
            hist            np.array    混淆矩阵，记录真实结果和预测结果，便于计算mIoUs。对于一个测试集而言，是以一张一张图片样本的方式，计算累加得到整个测试集的混淆矩阵
    '''
    print('Num classes', num_classes)  

    #-----------------------------------------#
    #   创建一个全是0的矩阵，是一个混淆矩阵
    #-----------------------------------------#
    hist = np.zeros((num_classes, num_classes))
    
    #------------------------------------------------#
    #   获得验证集标签路径列表，方便直接读取
    #   获得验证集图像分割结果路径列表，方便直接读取
    #------------------------------------------------#
    gt_imgs = [join(gt_dir, x + ".png") for x in png_name_list]  
    pred_imgs = [join(pred_dir, x + ".png") for x in png_name_list]  

    #------------------------------------------------#
    #   读取每一个（图片-标签）对
    #------------------------------------------------#
    print('          {:8s}{:8s}'.format('  mIoU', '   mPA'))
    for ind in range(len(gt_imgs)): 
        #------------------------------------------------#
        #   读取一张图像分割结果，转化成numpy数组
        #------------------------------------------------#
        pred = np.array(Image.open(pred_imgs[ind]))  
        #------------------------------------------------#
        #   读取一张对应的标签，转化成numpy数组
        #------------------------------------------------#
        label = np.array(Image.open(gt_imgs[ind]))  

        # 如果图像分割结果与标签的大小不一样，这张图片就不计算
        if len(label.flatten()) != len(pred.flatten()):  
            print(
                'Skipping: len(gt) = {:d}, len(pred) = {:d}, {:s}, {:s}'.format(
                    len(label.flatten()), len(pred.flatten()), gt_imgs[ind],
                    pred_imgs[ind]))
            continue

        #------------------------------------------------#
        #   对一张图片计算21×21的hist矩阵，并累加
        #------------------------------------------------#
        hist += fast_hist(label.flatten(), pred.flatten(), num_classes)  
        # 每计算10张就输出一下目前已计算的图片中所有类别平均的mIoU值
        if ind > 0 and ind % 10 == 0:  
            print('{:3d} / {:d}: {:5.2f}{:8.2f}'.format(ind, len(gt_imgs),
                                                    100 * np.nanmean(per_class_iu(hist)),
                                                    100 * np.nanmean(per_class_PA(hist))))
    #------------------------------------------------#
    #   计算所有验证集图片的逐类别mIoU值
    #------------------------------------------------#
    mIoUs   = per_class_iu(hist)
    mPA     = per_class_PA(hist)
    #------------------------------------------------#
    #   逐类别输出一下mIoU值
    #------------------------------------------------#
    print('Stastic of calculation results:')
    print('    {:15s}{:8s}{:8s}'.format('__class__', '_mIoU_', '_mPA_   '))
    for ind_class in range(num_classes):
        print('    {:15s}{:>5.2f}{:>8.2f}'.format(name_classes[ind_class],
                                                  mIoUs[ind_class] * 100,
                                                  mPA[ind_class]   * 100))

    #-----------------------------------------------------------------#
    #   在所有验证集图像上求所有类别平均的mIoU值，计算时忽略NaN值
    #-----------------------------------------------------------------#
    print('--> {:15s}{:>5.2f}{:>8.2f}'.format('mean', 
                                              np.nanmean(mIoUs) * 100,
                                              np.nanmean(mPA) * 100))

    reports = []
    return mIoUs, reports


if __name__ == "__main__":
    gt_dir        = "D:/【AI】/Datasets/CV_ds/VOC2012/VOCtrainval_11-May-2012/VOCdevkit/VOC2012/SegmentationClass/"
    pred_dir      = "./img/evaluate_results/"
    png_name_list = open("./model_data/test.txt",'r').read().splitlines() 
    #------------------------------#
    #   分类个数+1
    #   2+1
    #------------------------------#
    num_classes = 21
    #--------------------------------------------#
    #   区分的种类，和json_to_dataset里面的一样
    #--------------------------------------------#
    name_classes = ["background","aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
    compute_mIoU(gt_dir, pred_dir, png_name_list, num_classes, name_classes)  # 执行计算mIoU的函数
